import math
import random
import tkinter as tk
from tkinter import *

bg_color = "#CCFFFF"
btn_color = "#CCFFFF"

"""
此项目对电脑屏幕的显示设置有一定要求;
显示设置 ——> 缩放与布局 ——> 更改文本与应用等项目的大小 (100%)
显示设置 ——> 缩放与布局 ——> 显示器与分辨率 (1920 X 1080)
"""

class Application(Frame):
    def __init__(self, master=None):
        """初始化函数"""
        Frame.__init__(self, master)
        Pack.config(self)
        self.createwidgets()

    def createwidgets(self):
        # 创建画布
        width, height = 1800, 300

        self.canvas = Canvas(self, width=width, height=height, bg="white")
        self.canvas.pack()

        # label的x1/x2
        self.x1 = 6.15
        self.x2 = 10.15
        # 右边界
        self.scaling_right = 6.0
        # 左边界
        self.scaling_left = 0.3
        # 小球直径
        self.ball_d = 0.6
        # 存放小球x坐标数组
        self.ball_x = []
        # 存放小球y坐标数组
        self.ball_y = []
        # 存放小球x轴方向速度数组
        self.ball_v_x = []
        # 存放小球y轴方向速度数组
        self.ball_v_y = []
        # 创建小球
        self.ball01 = self.canvas.create_oval("6.00i", "1.50i", "6.6i", "2.10i", fill="red")
        self.ball02 = self.canvas.create_oval("10.0i", "1.50i", "10.6i", "2.10i", fill='black')

        # 设置x,y坐标，以及x,y方向速度
        self.ball01_x = 6.3
        self.ball01_y = 1.8

        self.ball02_x = 10.3
        self.ball02_y = 1.8

        # 选择系统
        self.buttonlist()
        # self.canvas.pack(side=LEFT)

    # 定义btn、label、entry
    def buttonlist(self):
        # 创建m1、m2、v1、v2、v1'、v2'标签、输入框、单位
        root.V1 = Label(root, text="速度v1", font=("黑体", 18), bg=btn_color)
        root.V1.place(relx=0.65, y=400, width=80, height=50)
        self.ball01_v_x = DoubleVar()
        root.entry_V1 = Entry(root, font=('微软雅黑', 18), fg="blue", textvariable=self.ball01_v_x)
        root.entry_V1.place(relx=0.70, y=400, width=80, height=50)
        self.ball01_v_x.set(0)
        root.V1_unit = Label(root, text="m/s", font=("黑体", 25), bg=btn_color)
        root.V1_unit.place(relx=0.745, y=405)

        root.V2 = Label(root, text="速度v2", font=("黑体", 18), bg=btn_color)
        root.V2.place(relx=0.80, y=400, width=80, height=50)
        self.ball02_v_x = DoubleVar()
        root.entry_V2 = Entry(root, font=('微软雅黑', 18), fg="blue", textvariable=self.ball02_v_x)
        root.entry_V2.place(relx=0.85, y=400, width=80, height=50)
        self.ball02_v_x.set(0)
        root.V2_unit = Label(root, text="m/s", font=("黑体", 25), bg=btn_color)
        root.V2_unit.place(relx=0.895, y=405)

        root.M1 = Label(root, text="质量m1", font=("黑体", 18), bg=btn_color)
        root.M1.place(relx=0.65, y=490, width=80, height=50)
        global m1
        m1 = DoubleVar()
        root.entry_M1 = Entry(root, font=('微软雅黑', 18), fg="blue", textvariable=m1)
        root.entry_M1.place(relx=0.70, y=490, width=80, height=50)
        m1.set(1.0)
        root.M1_unit = Label(root, text="kg", font=("黑体", 25), bg=btn_color)
        root.M1_unit.place(relx=0.745, y=490)

        root.M2 = Label(root, text="质量m2", font=("黑体", 18), bg=btn_color)
        root.M2.place(relx=0.80, y=490, width=80, height=50)
        global m2
        m2 = DoubleVar()
        root.entry_M2 = Entry(root, font=('微软雅黑', 18), fg="blue", textvariable=m2)
        root.entry_M2.place(relx=0.85, y=490, width=80, height=50)
        m2.set(1.0)
        root.M2_unit = Label(root, text="kg", font=("黑体", 25), bg=btn_color)
        root.M2_unit.place(relx=0.895, y=490)

        root.V11 = Label(root, text="速度v1’", font=("黑体", 18), bg=btn_color)
        root.V11.place(relx=0.63, y=580, width=120, height=50)
        self.ball011_v_x = DoubleVar()
        root.entry_V11 = Entry(root, font=('微软雅黑', 18), bg='#CCFFFF', fg="black", textvariable=self.ball011_v_x)
        root.entry_V11.place(relx=0.70, y=580, width=80, height=50)
        self.ball011_v_x.set(0)
        root.V11_unit = Label(root, text="m/s", font=("黑体", 25), bg=btn_color)
        root.V11_unit.place(relx=0.745, y=580)

        root.V22 = Label(root, text="速度v2’", font=("黑体", 18), bg=btn_color)
        root.V22.place(relx=0.78, y=580, width=120, height=50)
        self.ball022_v_x = DoubleVar()
        root.entry_V11 = Entry(root, font=('微软雅黑', 18), bg='#CCFFFF', fg="black", textvariable=self.ball022_v_x)
        root.entry_V11.place(relx=0.85, y=580, width=80, height=50)
        self.ball022_v_x.set(0)
        root.V11_unit = Label(root, text="m/s", font=("黑体", 25), bg=btn_color)
        root.V11_unit.place(relx=0.895, y=580)

        # 创建模式按钮(弹性碰撞、非弹性碰撞、完全非弹性碰撞) 单选框按钮
        global num
        num = tk.IntVar()
        Radiobutton(root, text="弹性碰撞", fg="blue", font=('微软雅黑', '30', 'bold'), indicatoron=False, variable=num,value=0).place(relx=0.15, y=360, width=300, height=80)
        Radiobutton(root, text="非弹性碰撞", fg="blue", font=('微软雅黑', '30', 'bold'), indicatoron=False, variable=num, value=1).place(relx=0.15, y=510, width=300, height=80)
        Radiobutton(root, text="完全非弹性碰撞", fg="blue", font=('微软雅黑', '30', 'bold'), indicatoron=False, variable=num, value=2).place(relx=0.15, y=660, width=300, height=80)

        # 初始机械能、碰后机械能、能量损耗最大值、实际能量损耗、单位
        global e1, e2, e3, e4
        root.E1 = Label(root, text="初始机械能", font=("黑体", 30), bg=btn_color, anchor='w')
        root.E1.place(relx=0.35, y=340, width=240, height=70)
        e1 = DoubleVar()
        root.entry_E1 = Entry(root, font=('微软雅黑', 25), fg="black", textvariable=e1, bg=btn_color)
        root.entry_E1.place(relx=0.51, y=350, width=100, height=50)
        e1.set(0)
        root.E1_unit = Label(root, text="J", font=("黑体", 30), bg=btn_color)
        root.E1_unit.place(relx=0.570, y=350)

        root.E2 = Label(root, text="碰后机械能", font=("黑体", 30), bg=btn_color, anchor='w')
        root.E2.place(relx=0.35, y=460, width=240, height=70)
        e2 = DoubleVar()
        root.entry_E2 = Entry(root, font=('微软雅黑', 25), fg="black", textvariable=e2, bg=btn_color)
        root.entry_E2.place(relx=0.51, y=470, width=100, height=50)
        e2.set(0)
        root.E2_unit = Label(root, text="J", font=("黑体", 30), bg=btn_color)
        root.E2_unit.place(relx=0.570, y=470)

        root.E3 = Label(root, text="实际能量损耗", font=("黑体", 30), bg=btn_color, anchor='w')
        root.E3.place(relx=0.35, y=580, width=240, height=70)
        e3 = DoubleVar()
        root.entry_E3 = Entry(root, font=('微软雅黑', 25), fg="black", textvariable=e3, bg=btn_color)
        root.entry_E3.place(relx=0.51, y=590, width=100, height=50)
        e3.set(0)
        root.E3_unit = Label(root, text="J", font=("黑体", 30), bg=btn_color)
        root.E3_unit.place(relx=0.570, y=590)

        root.E4 = Label(root, text="能量损耗最大值", font=("黑体", 30), bg=btn_color, anchor='w')
        root.E4.place(relx=0.35, y=700, width=280, height=70)
        e4 = DoubleVar()
        root.entry_E4 = Entry(root, font=('微软雅黑', 25), fg="black", textvariable=e4, bg=btn_color)
        root.entry_E4.place(relx=0.51, y=710, width=100, height=50)
        e4.set(0)
        root.E4_unit = Label(root, text="J", font=("黑体", 30), bg=btn_color)
        root.E4_unit.place(relx=0.570, y=710)

        # 创建开始按钮
        root.btn01 = Button(root, text="确定", font=("黑体", 30), command=self.start, fg="red")
        root.btn01.place(relx=0.45, y=820, width=100, height=80)

        # 创建重新开始按钮
        root.btn02 = Button(root, text="重新开始", font=("黑体", 30), command=self.restart, fg="red")
        root.btn02.place(relx=0.60, y=820, width=200, height=80)

    def start(self):
        self.update_ball_x_y()
        e1.set(1 / 2 * (m1.get() * self.ball_v_x[0] * self.ball_v_x[0] + m2.get() * self.ball_v_x[1] * self.ball_v_x[1]))
        e4.set(format(e1.get() - 1 / 2 * (m1.get() + m2.get()) * pow((m1.get() * self.ball_v_x[0] + m2.get() * self.ball_v_x[1]) / (m1.get() + m2.get()), 2), '.2f'))
        self.moveball()

    def restart(self):
        self.canvas.destroy()
        self.createwidgets()

    # 更新小球的信息
    def update_ball_x_y(self):
        # 第一个小球信息
        self.ball_x.append(self.ball01_x)
        self.ball_y.append(self.ball01_y)
        self.ball_v_x.append(self.ball01_v_x.get())
        # 第二个小球信息
        self.ball_x.append(self.ball02_x)
        self.ball_y.append(self.ball02_y)
        self.ball_v_x.append(self.ball02_v_x.get())

    # 小球移动
    def moveball(self):
        if num.get() == 0:
            self.update_ball_v_elastic_collision()
        if num.get() == 1:
            self.update_ball_v_inelastic_collisions()
        if num.get() == 2:
            self.update_ball_v_completely_inelastic_collision()
        deltax = self.get_ball_deltax(0)
        self.x1 = deltax + self.x1
        self.label_ball1 = Label(self, text="m1", bg="red", font=("黑体", 20))
        self.label_ball1.place(x="%ri" % self.x1, y="1.6i")
        self.label_ball1.after(10, lambda: self.label_ball1.destroy())
        self.canvas.move(self.ball01, "%ri" % deltax, 0)

        deltax = self.get_ball_deltax(1)
        self.x2 = self.x2 + deltax
        self.label_ball2 = Label(self, text="m2", bg="black", font=("黑体", 20), fg="white")
        self.label_ball2.place(x="%ri" % self.x2, y="1.6i")
        self.label_ball2.after(10, lambda: self.label_ball2.destroy())
        self.canvas.move(self.ball02, "%ri" % deltax, 0)

        self.after(10, self.moveball)

    # 更新小球的速度（弹性碰撞）
    def update_ball_v_elastic_collision(self):
        # 小球碰撞条件，即：(x2 - x1)^2 + (y2 - y1)^2 <= (r + R)^2
        if (round((self.ball_x[0] - self.ball_x[1]) ** 2 + (self.ball_y[0] - self.ball_y[1]) ** 2, 2) <= round(self.ball_d ** 2, 2)):
            # 两小球速度交换
            temp_vx1 = self.ball_v_x[0]
            temp_vx2 = self.ball_v_x[1]
            self.ball_v_x[0] = ((m1.get() - m2.get()) * self.ball_v_x[0] + 2 * m2.get() * self.ball_v_x[1]) / (m1.get() + m2.get())
            self.ball_v_x[1] = ((m2.get() - m1.get()) * temp_vx2 + 2 * m1.get() * temp_vx1) / (m1.get() + m2.get())
            e2.set(1 / 2 * (m1.get() * self.ball_v_x[0] * self.ball_v_x[0] + m2.get() * self.ball_v_x[1] * self.ball_v_x[1]))
            e3.set(format(e1.get() - e2.get(), '.3f'))
            self.ball011_v_x.set(format(self.ball_v_x[0], '.3f'))
            self.ball022_v_x.set(format(self.ball_v_x[1], '.3f'))

    # 更新小球的速度（非弹性碰撞）
    def update_ball_v_inelastic_collisions(self):
        if (round((self.ball_x[0] - self.ball_x[1]) ** 2 + (self.ball_y[0] - self.ball_y[1]) ** 2, 2) <= round(self.ball_d ** 2, 2)):
            e3.set(format(random.uniform(0, e4.get()), '.2f'))
            temp_vx1 = self.ball_v_x[0]
            temp_vx2 = self.ball_v_x[1]
            self.ball_v_x[0] = (m1.get() * temp_vx1 + m2.get() * temp_vx2 - m2.get() * ((m1.get() * temp_vx1 + m2.get() * temp_vx2) / (m1.get() + m2.get()) + math.sqrt(m1.get() * m2.get() * (-2.0 * e3.get() * m1.get() - 2.0 * e3.get() * m2.get() + m1.get() * m2.get() * temp_vx1 ** 2 - 2.0 * m1.get() * m2.get() * temp_vx1 * temp_vx2 + m1.get() * m2.get() * temp_vx2 ** 2)) / (m2.get() * (m1.get() + m2.get())))) / m1.get()
            self.ball_v_x[1] = (m1.get() * temp_vx1 + m2.get() * temp_vx2) / (m1.get() + m2.get()) + math.sqrt(m1.get() * m2.get() * (-2.0 * e3.get() * m1.get() - 2.0 * e3.get() * m2.get() + m1.get() * m2.get() * temp_vx1 ** 2 - 2.0 * m1.get() * m2.get() * temp_vx1 * temp_vx2 + m1.get() * m2.get() * temp_vx2 ** 2)) / (m2.get() * (m1.get() + m2.get()))
            e2.set(format(1 / 2 * (m1.get() * self.ball_v_x[0] * self.ball_v_x[0] + m2.get() * self.ball_v_x[1] * self.ball_v_x[1]), '.3f'))
            self.ball011_v_x.set(format(self.ball_v_x[0], '.3f'))
            self.ball022_v_x.set(format(self.ball_v_x[1], '.3f'))

    # 更新小球的速度（完全非弹性碰撞）
    def update_ball_v_completely_inelastic_collision(self):
        # 小球碰撞条件，即：(x2 - x1)^2 + (y2 - y1)^2 <= (r + R)^2
        if (round((self.ball_x[0] - self.ball_x[1]) ** 2 + (self.ball_y[0] - self.ball_y[1]) ** 2, 2) <= round(self.ball_d ** 2, 2)):
            # 两小球黏在一起运动
            self.ball_v_x[0] = (m1.get() * self.ball_v_x[0] + m2.get() * self.ball_v_x[1]) / (m1.get() + m2.get())
            self.ball_v_x[1] = self.ball_v_x[0]
            e2.set(1 / 2 * (m1.get() * self.ball_v_x[0] * self.ball_v_x[0] + m2.get() * self.ball_v_x[1] * self.ball_v_x[1]))
            e3.set(e1.get() - e2.get())
            self.ball011_v_x.set(format(self.ball_v_x[0], '.3f'))
            self.ball022_v_x.set(format(self.ball_v_x[1], '.3f'))

    # 获取小球的x增量
    def get_ball_deltax(self, index):
        '''获取小球X轴坐标移动距离并且更新小球的圆心X坐标，返回X轴所需移动距离'''
        deltax = self.ball_v_x[index] / 100
        self.ball_x[index] = self.ball_x[index] + deltax
        return deltax


if __name__ == '__main__':
    root = Tk()
    root.title("两球碰撞仿真实验")  # 改标题
    root.config(background=bg_color)  # 改变背景颜色
    screen_height = root.winfo_screenheight()
    screen_width = root.winfo_screenwidth()
    root.geometry("%sx%s" % (screen_width, screen_height))  # 设置宽、高
    Application(master=root)
    root.mainloop()
